// rvv_agg.c
#include <jni.h>
#include <stdlib.h>
#include <string.h>
#include <riscv_vector.h>

// 聚合状态结构
typedef struct {
    double sum_qty;
    double sum_price;
    double sum_disc_price;
    double sum_charge;
    double sum_disc;
    long count;
} AggState;

// 创建并初始化聚合状态
JNIEXPORT jlong JNICALL Java_TPCHRVVAggregator_createAggState(JNIEnv *env, jobject obj) {
    AggState* state = (AggState*)malloc(sizeof(AggState));
    if (state == NULL) {
        return 0;
    }
    memset(state, 0, sizeof(AggState));
    return (jlong)state;
}

// 使用RVV向量化指令批量更新聚合状态
JNIEXPORT void JNICALL Java_TPCHRVVAggregator_updateBatch(
    JNIEnv *env, jobject obj,
    jlong statePtr,
    jdoubleArray quantitiesArr,
    jdoubleArray pricesArr,
    jdoubleArray discountsArr,
    jdoubleArray taxesArr,
    jint count)
{
    AggState* state = (AggState*)statePtr;
    jdouble *quantities = (*env)->GetDoubleArrayElements(env, quantitiesArr, NULL);
    jdouble *prices = (*env)->GetDoubleArrayElements(env, pricesArr, NULL);
    jdouble *discounts = (*env)->GetDoubleArrayElements(env, discountsArr, NULL);
    jdouble *taxes = (*env)->GetDoubleArrayElements(env, taxesArr, NULL);

    size_t vl;
    size_t n = count;

    // 标量累加器
    double batch_sum_qty = 0.0;
    double batch_sum_price = 0.0;
    double batch_sum_disc_price = 0.0;
    double batch_sum_charge = 0.0;
    double batch_sum_disc = 0.0;

    for (size_t i = 0; i < n; i += vl) {
        // 修正3.2.1: 使用正确的vsetvl函数名
        vl = __riscv_vsetvl_e64m1(n - i);

        // 加载向量
        vfloat64m1_t v_qty = __riscv_vle64_v_f64m1(&quantities[i], vl);
        vfloat64m1_t v_price = __riscv_vle64_v_f64m1(&prices[i], vl);
        vfloat64m1_t v_discount = __riscv_vle64_v_f64m1(&discounts[i], vl);
        vfloat64m1_t v_tax = __riscv_vle64_v_f64m1(&taxes[i], vl);

        // 计算中间值
        vfloat64m1_t v_one_minus_disc = __riscv_vfrsub_vf_f64m1(v_discount, 1.0, vl);
        vfloat64m1_t v_disc_price = __riscv_vfmul_vv_f64m1(v_price, v_one_minus_disc, vl);
        vfloat64m1_t v_one_plus_tax = __riscv_vfadd_vf_f64m1(v_tax, 1.0, vl);
        vfloat64m1_t v_charge = __riscv_vfmul_vv_f64m1(v_disc_price, v_one_plus_tax, vl);

        // 为归约操作创建初始值向量
        vfloat64m1_t v_red_init = __riscv_vfmv_v_f_f64m1(0.0, vl);

        // 修正3.2.2: 使用正确的归约函数名和参数
        vfloat64m1_t v_red_qty = __riscv_vfredusum_vs_f64m1_f64m1(v_qty, v_red_init, vl);
        batch_sum_qty += __riscv_vfmv_f_s_f64m1_f64(v_red_qty);

        vfloat64m1_t v_red_price = __riscv_vfredusum_vs_f64m1_f64m1(v_price, v_red_init, vl);
        batch_sum_price += __riscv_vfmv_f_s_f64m1_f64(v_red_price);

        vfloat64m1_t v_red_disc_price = __riscv_vfredusum_vs_f64m1_f64m1(v_disc_price, v_red_init, vl);
        batch_sum_disc_price += __riscv_vfmv_f_s_f64m1_f64(v_red_disc_price);

        vfloat64m1_t v_red_charge = __riscv_vfredusum_vs_f64m1_f64m1(v_charge, v_red_init, vl);
        batch_sum_charge += __riscv_vfmv_f_s_f64m1_f64(v_red_charge);

        vfloat64m1_t v_red_discount = __riscv_vfredusum_vs_f64m1_f64m1(v_discount, v_red_init, vl);
        batch_sum_disc += __riscv_vfmv_f_s_f64m1_f64(v_red_discount);
    }

    // 累加到状态
    state->sum_qty += batch_sum_qty;
    state->sum_price += batch_sum_price;
    state->sum_disc_price += batch_sum_disc_price;
    state->sum_charge += batch_sum_charge;
    state->sum_disc += batch_sum_disc;
    state->count += n;

    // 释放数组
    (*env)->ReleaseDoubleArrayElements(env, quantitiesArr, quantities, JNI_ABORT);
    (*env)->ReleaseDoubleArrayElements(env, pricesArr, prices, JNI_ABORT);
    (*env)->ReleaseDoubleArrayElements(env, discountsArr, discounts, JNI_ABORT);
    (*env)->ReleaseDoubleArrayElements(env, taxesArr, taxes, JNI_ABORT);
}

// 合并两个聚合状态
JNIEXPORT void JNICALL Java_TPCHRVVAggregator_merge(JNIEnv *env, jobject obj, jlong statePtr1, jlong statePtr2) {
    AggState* s1 = (AggState*)statePtr1;
    AggState* s2 = (AggState*)statePtr2;

    s1->sum_qty += s2->sum_qty;
    s1->sum_price += s2->sum_price;
    s1->sum_disc_price += s2->sum_disc_price;
    s1->sum_charge += s2->sum_charge;
    s1->sum_disc += s2->sum_disc;
    s1->count += s2->count;
}

// 获取最终聚合结果
JNIEXPORT jobjectArray JNICALL Java_TPCHRVVAggregator_getResult(JNIEnv *env, jobject obj, jlong statePtr) {
    AggState* state = (AggState*)statePtr;

    jclass doubleClass = (*env)->FindClass(env, "java/lang/Double");
    if (doubleClass == NULL) return NULL; // 类未找到

    jmethodID doubleConstructor = (*env)->GetMethodID(env, doubleClass, "<init>", "(D)V");
    if (doubleConstructor == NULL) return NULL; // 构造函数未找到

    jobjectArray result = (*env)->NewObjectArray(env, 7, doubleClass, NULL);
    if (result == NULL) return NULL; // 数组创建失败

    // 辅助函数设置数组元素
    int index = 0;
    #define SET_DOUBLE_IN_ARRAY(value) do { \
        jobject d_obj = (*env)->NewObject(env, doubleClass, doubleConstructor, (value)); \
        (*env)->SetObjectArrayElement(env, result, index++, d_obj); \
        (*env)->DeleteLocalRef(env, d_obj); \
    } while(0)
    
    double avg_qty = (state->count == 0) ? 0.0 : state->sum_qty / state->count;
    double avg_price = (state->count == 0) ? 0.0 : state->sum_price / state->count;
    double avg_disc = (state->count == 0) ? 0.0 : state->sum_disc / state->count;

    SET_DOUBLE_IN_ARRAY(state->sum_qty);
    SET_DOUBLE_IN_ARRAY(state->sum_price);
    SET_DOUBLE_IN_ARRAY(state->sum_disc_price);
    SET_DOUBLE_IN_ARRAY(state->sum_charge);
    SET_DOUBLE_IN_ARRAY(avg_qty);
    SET_DOUBLE_IN_ARRAY(avg_price);
    SET_DOUBLE_IN_ARRAY(avg_disc);

    #undef SET_DOUBLE_IN_ARRAY

    return result;
}

// 释放聚合状态占用的内存
JNIEXPORT void JNICALL Java_TPCHRVVAggregator_freeState(JNIEnv *env, jobject obj, jlong statePtr) {
    free((void*)statePtr);
}
